import numpy as np


def dcm_from_euler(euler_angles):
    """
    计算方向余弦矩阵（DCM）(Earth2Body）
    :param roll: 滚转角，单位为弧度
    :param pitch: 俯仰角，单位为弧度
    :param yaw: 偏航角，单位为弧度
    :return: 方向余弦矩阵（DCM）
    """
    roll, pitch, yaw = euler_angles
    c_roll = np.cos(roll)
    s_roll = np.sin(roll)
    c_pitch = np.cos(pitch)
    s_pitch = np.sin(pitch)
    c_yaw = np.cos(yaw)
    s_yaw = np.sin(yaw)

    # 构建方向余弦矩阵
    dcm = np.array([
        [c_yaw * c_pitch, c_yaw * s_pitch * s_roll - s_yaw * c_roll, c_yaw * s_pitch * c_roll + s_yaw * s_roll],
        [s_yaw * c_pitch, s_yaw * s_pitch * s_roll + c_yaw * c_roll, s_yaw * s_pitch * c_roll - c_yaw * s_roll],
        [-s_pitch, c_pitch * s_roll, c_pitch * c_roll]
    ])
    return dcm


def rotation_matrix(euler_angles):  # （Body2Earth)
    phi, theta, psi = euler_angles
    R_x = np.array([[1, 0, 0],
                    [0, np.cos(phi), -np.sin(phi)],
                    [0, np.sin(phi), np.cos(phi)]])

    R_y = np.array([[np.cos(theta), 0, np.sin(theta)],
                    [0, 1, 0],
                    [-np.sin(theta), 0, np.cos(theta)]])

    R_z = np.array([[np.cos(psi), -np.sin(psi), 0],
                    [np.sin(psi), np.cos(psi), 0],
                    [0, 0, 1]])

    return R_z @ R_y @ R_x


def angle_body2world(euler_angles):
    phi, theta, psi = euler_angles
    Trans_matrix = np.array([[1, np.tan(theta) * np.sin(phi), np.tan(theta) * np.cos(phi)],
                             [0, np.cos(phi), -np.sin(phi)],
                             [0, np.sin(phi) / np.cos(theta), np.cos(phi) / np.cos(theta)]])
    return Trans_matrix


def rk4_step(f, state, u, t, dt):
    k1 = f(state, u, t)
    print([x * 0.5 * dt for x in k1])
    k2 = f(state + [x * 0.5 * dt for x in k1], u, t + 0.5 * dt)
    k3 = f(state + [x * 0.5 * dt for x in k2], u, t + 0.5 * dt)
    k4 = f(state + [x  * dt for x in k3], u, t + dt)

    return state + (dt / 6) * (k1 + 2 * k2 + 2 * k3 + k4)
